# -*- coding:utf-8 -*-


'''
验证码图片数据集及相关图片处理方法
'''


import string
from numpy import *
import cv2
import os
import time
import requests
from PIL import Image
import time


# 获取验证码图片
def getPic(url, name):
    name = str(name)
    imgresponse = requests.get(url, stream=True)
    image = imgresponse.content
    folder = '../img2' + os.path.sep
    print('save img: ' + folder + name + '.jpg')
    try:
        with open(folder + name + '.jpg', "wb") as jpg:
            jpg.write(image)
            jpg.flush()
    except IOError:
        print("IO Error\n")
    # 将图片转一次格式，否则opencv读取不到图片
    im = Image.open(folder + name + '.jpg')
    im.save(folder + name + '.png')
    os.remove(folder + name + '.jpg')

# 图片去噪
def removeSalt(img, fileName):
    img = img[1:26,5:55]
    median = cv2.medianBlur(img, 3)
    print('img: '  + str(fileName) + ' remove salt')
    return median

# 二值化
def binary(img, fileName):
    ret, thresh1 = cv2.threshold(img, 108, 255, cv2.THRESH_BINARY)
    print('img: ' + str(fileName) + ' binary')
    return thresh1

# 分割图片
def cutPic(img, fileName):
    folder = '../newimg' + os.path.sep
    n = shape(img)[1]
    width = int(n / 4)
    arr = []
    for i in range(4):
        sonImg = img[:,i * width: (i + 1) * width]
        sonFileName = folder + fileName[:fileName.rfind('.')] + str(i) + ".png"
        cv2.imwrite(sonFileName, sonImg)
        arr.append(sonFileName)
    print('img: '  + str(fileName) + ' cut')
    return arr

# 把图片转化为行向量
def img2Vec(imgPath):
    img = cv2.imread(imgPath,0)
    returnVec = zeros(shape(img), dtype = float32)
    returnVec[img == 0] = 1
    returnVec[img == 255] = 0
    returnVec = returnVec.flatten()
    return returnVec

# 去噪、二值化、切割图片，按顺序返回所有子图路径
def getAllSonImg(imgPath):
    img = cv2.imread(imgPath, 0)
    index = imgPath.rfind(os.path.sep)
    if index == -1:
        index = imgPath.rfind('/')
    fileName = imgPath[index + 1:]
    img1 = removeSalt(img, fileName)
    img2 = binary(img1, fileName)
    return cutPic(img2, fileName)


# 获取标签与序号对应的字典
def getLabelIndexDir():
    dir = {}
    i = 0
    for j in range(9):
        dir[str(j)] = i
        i = i + 1
    for char in string.ascii_lowercase:
        if char == 'o' or char == 'z':
            continue
        dir[char] = i
        i = i + 1
    return dir

# 加载训练集
def _loadDataSet(dirPath):
    trainSet = []
    labels = []
    labelIndexDir = getLabelIndexDir()
    pics = os.listdir(dirPath)
    for pic in pics:
        label = pic[0]
        vec = img2Vec(dirPath + os.path.sep + pic)
        trainSet.append(vec)
        labelVec = zeros(33)
        labelVec[labelIndexDir[label]] = 1.0
        labels.append(labelVec)
    return mat(trainSet), mat(labels)

import base64

# 将图片编码为base64字符串
def img2base64(imgPath):
    f = open(imgPath, 'rb')  # 二进制方式打开图文件
    ls_f = base64.b64encode(f.read())  # 读取文件内容，转换为base64编码
    f.close()
    return base64.urlsafe_b64encode(ls_f).rstrip('='.encode())

# 将base64编码的字符串解码为图片
def base642img(base64Str):
    path = "../base64img"
    if(type(base64Str) == 'bytes'):
        addStr = '='.encode()
    else:
        addStr = '='
    try:
        base64Str = base64.urlsafe_b64decode(str(base64Str + addStr * (4 - len(base64Str) % 4)))
    except Exception as ex:
        return None
    imgdata = base64.b64decode(base64Str)
    fileName = os.path.join(path, str(time.time()) + '.jpg')
    file = open(fileName, 'wb')
    file.write(imgdata)
    file.close()
    return fileName

'''
数据集类
'''
class VertifyImg(object):
    def __init__(self, trainImgDir, testImgDir):
        self.trainDataSet,self.trainLabels = _loadDataSet(trainImgDir)
        self.testDataSet, self.testLabels = _loadDataSet(testImgDir)
        self.trainSize = len(self.trainDataSet)
        self.testSize = len(self.testDataSet)

    def next_batch(self, batch_size):
        if batch_size > self.trainSize:
            print("There is not enough data!")
            return None
        start_index = random.randint(0, self.trainSize - batch_size)
        return self.trainDataSet[start_index:(start_index + batch_size)], \
               self.trainLabels[start_index:(start_index + batch_size)]

    def get_test_data(self):
        return self.testDataSet, self.testLabels



